#include "yaml-cpp/yaml.h"
#include <iostream>
#include <iostream>
#include "db_conninfo.h"
#include "opengauss/libpq-fe.h"
#include <iomanip> // 包含setw
#include <fstream>

std::string getsource(const char *location)
{
    YAML::Node config = YAML::LoadFile(location);
    std::string datasource;
    datasource.append("dbname=");
    datasource.append(config["diff"]["source"]["database"].as<std::string>());
    datasource.append(" port=");
    datasource.append(config["diff"]["source"]["port"].as<std::string>());
    datasource.append(" host=");
    datasource.append(config["diff"]["source"]["host"].as<std::string>());
    datasource.append(" application_name=libpq connect_timeout=5 sslmode=allow");
    datasource.append(" user=");
    datasource.append(config["diff"]["source"]["username"].as<std::string>());
    if (!config["diff"]["source"]["password"].IsNull())
    {
        datasource.append(" password=");
        datasource.append("'");
        datasource.append(config["diff"]["source"]["password"].as<std::string>());
        datasource.append("'");
    }
    return datasource;
}

static void
exit_nicely(PGconn *conn)
{
    PQfinish(conn);
    exit(1);
}

int QuerySettings(PGconn *conn, const char *sql, std::vector<std::vector<std::string>> &header, std::vector<std::vector<std::string>> &data)
{
    int nFields;
    int i, j;

    // 执行查询语句
    PGresult *res = PQexec(conn, sql);
    if (PQresultStatus(res) != PGRES_TUPLES_OK)
    {
        fprintf(stderr, "查询失败: %s", PQerrorMessage(conn));
        PQclear(res);
        PQfinish(conn);
        return 1;
    }
    /* 打印属性名称 */
    std::cout << std::left; // 设置对齐方式为左对齐
    nFields = PQnfields(res);
    std::vector<std::string> headerdata;
    for (i = 0; i < nFields; i++)
    {
        headerdata.push_back(PQfname(res, i));
        std::cout << std::setw(25) << PQfname(res, i);
    }
    header.push_back(headerdata);
    std::cout << std::endl;
    /* 打印行 */
    for (i = 0; i < PQntuples(res); i++)
    {
        std::vector<std::string> rowData;
        for (j = 0; j < nFields; j++)
        {
            rowData.push_back(PQgetvalue(res, i, j));
            std::cout << std::setw(25) << PQgetvalue(res, i, j);
        }
        data.push_back(rowData);
        std::cout << std::endl;
    }
    // 释放结果集和关闭数据库连接
    PQclear(res);
    return 0;
}

void writeToCsv(const std::string &filename, const std::vector<std::vector<std::string>> &data)
{
    std::ofstream file(filename, std::ios::app);
    if (!file.is_open())
    {
        throw std::runtime_error("无法打开文件: " + filename);
    }

    for (const auto &row : data)
    {
        for (size_t i = 0; i < row.size(); ++i)
        {
            // 处理包含逗号或引号的值
            std::string value = row[i];
            bool needsQuotes = (value.find(',') != std::string::npos ||
                                value.find('"') != std::string::npos ||
                                value.find('\n') != std::string::npos);

            if (needsQuotes)
            {
                file << '"';
                // 转义引号
                for (char c : value)
                {
                    if (c == '"')
                        file << '"';
                    file << c;
                }
                file << '"';
            }
            else
            {
                file << value;
            }

            if (i != row.size() - 1)
            {
                file << ','; // 分隔符
            }
        }
        file << '\n'; // 换行符
    }

    file.close();
}

int main(int argc, char **argv)
{
    const char *location;
    if (argc > 1)
        location = argv[1];
    else
        location = "inspection.yaml";
    std::string srouce;
    srouce = getsource(location);
    PGconn *conn = createPGConnection(srouce);
    if (!conn)
    {
        return EXIT_FAILURE;
    }
    int status;
    YAML::Node config = YAML::LoadFile(location);
    for (YAML::const_iterator it = config["diff"]["privilege_type"].begin(); it != config["diff"]["privilege_type"].end(); ++it)
    {
        std::cout << it->first.as<std::string>() << std::endl;
        std::vector<std::vector<std::string>> header;
        std::vector<std::vector<std::string>> tableData;
        status = QuerySettings(conn, it->second.as<std::string>().c_str(),header,tableData);
        if (status != 0)
        {
            std::cerr << "连接数据库出错: " << PQerrorMessage(conn) << std::endl;
            return EXIT_FAILURE;
        }
        writeToCsv("output.csv", header);
        writeToCsv("output.csv", tableData);
    }
    PQfinish(conn);

    return 0;
}
